# 一个点云和一张图片组成一个 sample, 点云数据用 npz 文件存
# dots 字段存了点云数据, shape 为 (n, c, r, 2)
# id 字段存了每个点云对应的图片编号, 也就是文件名 0001.jpg

import numpy as np
import cv2
from tqdm import tqdm
from pathlib import Path
import shutil
BASE_DIR = Path(__file__).resolve().parent

from PyQt5 import QtGui, QtWidgets
from PyQt5.QtCore import pyqtSignal, QTimer, Qt
from gelslimpi.visualize import QImageViewWidget, QTablePanel, QCheckList
from gelslimpi.utils import now

class MarkerData():
    def __init__(self, data_path):
        self.data_path = data_path
        data = np.load(self.data_path / "data.npz")
        self.dots = data["dots"]
        if "id" in data.files:
            self.ids = data["id"]
        else:
            self.ids = np.arange(len(self.dots))

        self.row, self.col = self.dots.shape[1:3]

        self.curr_id = 0

    @property
    def len(self):
        return len(self.dots)

    def get_sample(self, idx=None):
        idx = self.curr_id if idx is None else idx % self.len 
        img_path = self.data_path / f"{self.ids[idx]:04d}.jpg"
        img = cv2.imread(str(img_path))
        return img, self.dots[idx], img_path

    @staticmethod
    def draw_marker(img, dots):
        img = img.copy()
        row, col = dots.shape[:2]
        dots = dots.reshape(-1, 2)
        for i, dot in enumerate(dots):
            #color 为hsv渐变色
            color = np.array([[[i / col * 180/7, 255, 255]]], dtype=np.uint8)
            color = cv2.cvtColor(color, cv2.COLOR_HSV2RGB)[0][0]
            cv2.circle(img, (int(dot[0]), int(dot[1])), 10, tuple(int(x) for x in color), 1)
        return img

    def del_samples(self, del_ids):
        if isinstance(del_ids, int):
            del_ids = [del_ids]

        del_ids = [i % self.len for i in del_ids]

        for del_id in del_ids:
            img_path = self.data_path / f"{self.ids[del_id]:04d}.jpg"
            img_path.unlink(missing_ok=True)
            print(img_path, "deleted")
        self.dots = np.delete(self.dots, del_ids, axis=0)
        self.ids = np.delete(self.ids, del_ids, axis=0)
        self._len = len(self.dots)

    def save(self):
        print("Saved!")
        np.savez(self.data_path / "data.npz", dots=self.dots, id=self.ids)

    def search_dot(self, x=None, y=None, idx=None):
        idx = self.curr_id if idx is None else idx % self.len 
        dots = self.dots[idx]
        dots = dots.reshape(-1, 2)
        dist = np.sum((dots - np.array([x, y]))**2, axis=1)
        idx = np.argmin(dist.copy())
        return (idx//self.col, idx%self.col)

    def update_dot(self, row=None, col=None, x=None, y=None, idx=None):
        idx = self.curr_id if idx is None else idx % self.len 
        self.dots[idx, row, col] = np.array([x, y])

    def next(self):
        self.curr_id = (self.curr_id+1) % self.len
    
    def prev(self):
        self.curr_id = (self.curr_id-1+self.len) % self.len

    def _del(self):
        self.del_samples(self.curr_id)
        self.curr_id = (self.curr_id) % self.len


class DataClean(QtWidgets.QWidget):

    pannel_config = {
        "auto_save": ["bool", 1],
        "output_dir": ["directory", ""],
        "merge + export": ["button", ], 
        "spacer0": ["spacer", 10],
        "line1": ["line"],
        "spacer1": ["spacer", 10],
        "del sample": ["button", ],
        "save current dataset": ["button", ],
        "spacer2": ["spacer", 20],
        "line2": ["line"],
        "spacer3": ["spacer", 20],
        "prev": ["button", ],
        "next": ["button", ],
        "spacer4": ["spacer", 20],
        "line3": ["line"],
        "spacer5": ["spacer", 20],
    }

    def __init__(
        self, 
        data_dir = BASE_DIR / "runs",
        output_dir = BASE_DIR / ("runs/output-" + now("%m%d_%H%M%S")),
    ):
        super().__init__()
        self.data_dir = Path(data_dir)
        self.datasets = []
        self.dirs = self._load_dirs(self.data_dir)
        self.output_dir = Path(output_dir)
        self.pannel_config["output_dir"][1] = str(output_dir)
        self.init_ui()

        self.timer = QTimer(self)
        self.timer.timeout.connect(self.on_timeout)
        self.timer.start(50)

        self.pannel.sigTableChanged.connect(self.on_table_changed)
        self.img_view.sigMousePressed.connect(self.on_mouse_pressed)
        self.img_view.sigMousePressedMove.connect(self.on_mouse_pressed_move)

    @property
    def curr_dataset(self):
        selected = self.pannel.widget("dirs").value
        if len(selected) == 0:
            print("selected 0")
            self.pannel["dirs"] = 0
            return self.datasets[0]
        else:
            return self.datasets[selected[0]]

    def init_ui(self):
        self.setWindowTitle("Data Clean")
        self.resize(1480, 960)

        hbox = QtWidgets.QHBoxLayout(self)
        hbox.setContentsMargins(0, 0, 0, 0)
        hbox.setSpacing(0)

        left_frame = QtWidgets.QFrame(self)
        line = QtWidgets.QFrame(self)
        line.setFrameShape(QtWidgets.QFrame.VLine)
        line.setFrameShadow(QtWidgets.QFrame.Sunken)

        self.pannel_config["dirs"] = ["checklist", 0, self.dirs, "v", 0]
        self.pannel = QTablePanel(self, self.pannel_config)
        hbox.addWidget(left_frame, 12)
        hbox.addWidget(line)
        hbox.addWidget(self.pannel, 3)

        vbox = QtWidgets.QVBoxLayout(left_frame)
        vbox.setContentsMargins(0, 0, 0, 0)
        vbox.setSpacing(5)
        self.img_view = QImageViewWidget(left_frame, auto_scale=True)
        self.status_bar = QtWidgets.QLabel(left_frame)
        vbox.addWidget(self.img_view, 15)
        vbox.addWidget(self.status_bar, 1)

    def _load_dirs(self, base_dir: Path):
        dirs = []
        for dir in base_dir.glob("*"):
            if not dir.is_dir() or not (dir / "data.npz").exists():
                continue
            dirs.append(dir)
            self.datasets.append(MarkerData(dir))
        return dirs

    def on_timeout(self):
        img, dots, img_path = self.curr_dataset.get_sample()
        img = MarkerData.draw_marker(img, dots)
        self.img_view.setData(img)
        self.status_bar.setText(f"idx: {self.curr_dataset.curr_id:04d}/{self.curr_dataset.len} | path: {img_path}")

    def on_table_changed(self, param_key):
        param_val = self.pannel[param_key]
        if param_key == "prev":
            self.curr_dataset.prev()
        elif param_key == "next":
            self.curr_dataset.next()
        elif param_key == "del":
            self.curr_dataset._del()
        elif param_key == "save":
            self.curr_dataset.save()
        elif param_key == "merge + export":
            self.merge_export()

    def keyPressEvent(self, a0) -> None:
        if a0.key() in [Qt.Key.Key_Left, Qt.Key.Key_A]:
            self.curr_dataset.prev()
        elif a0.key() in [Qt.Key.Key_Right, Qt.Key.Key_D]:
            self.curr_dataset.next()
        elif a0.key() in [Qt.Key.Key_X,]:
            self.curr_dataset._del()
        return super().keyPressEvent(a0)

    def on_mouse_pressed(self, xy):
        x, y = xy
        self.row, self.col = self.curr_dataset.search_dot(x, y)
        self.curr_dataset.update_dot(self.row, self.col, x, y)
        print(f"row: {self.row}, col: {self.col}, x: {x}, y: {y}")

    def on_mouse_pressed_move(self, xy):
        x, y = xy
        self.curr_dataset.update_dot(self.row, self.col, x, y)
        print(f"row: {self.row}, col: {self.col}, x: {x}, y: {y}")

    def closeEvent(self, a0) -> None:
        if self.pannel["auto_save"]:
            for dataset in self.datasets:
                dataset.save()
        return super().closeEvent(a0)

    def merge_export(self):
        selected = self.pannel.widget("dirs").value
        if len(selected) == 0:
            print("no dataset selected")
            return
        self.output_dir = Path(self.pannel["output_dir"])
        self.output_dir.mkdir(exist_ok=True, parents=True)
        dots = []
        cnt = 0
        for i in tqdm(selected, total=len(selected)):
            dataset = self.datasets[i]
            dots.append(dataset.dots)
            for id in tqdm(dataset.ids, total=dataset.len, leave=False):
                src_path = dataset.data_path / f"{id:04d}.jpg"
                dst_path = self.output_dir / f"{cnt:04d}.jpg"
                cnt += 1
                shutil.copy(src_path, dst_path)
        dots = np.concatenate(dots, axis=0)
        np.savez(self.output_dir / "data.npz", dots=dots, id=np.arange(cnt))
        print("export successfully")

if __name__ == '__main__':
    import sys
    data_dir = "/home/mpi/workspace/gelslim-calibrate/data/marker"
    # data_dir = BASE_DIR / "runs/"

    app = QtWidgets.QApplication(sys.argv)
    win = DataClean(
        data_dir = data_dir
    )
    win.show()
    sys.exit(app.exec_())